# encoding: utf-8
# @Time    : 2024/08/13
# @Author  : 半只程序员
# @Email   : 18152007693@163.com
# @File    : compare_row_plus.py
# @Software: PyCharm
"""
# 两个文件数据比较，并输出对比结果到文件，不同的数据会以红色字体显示
"""
import os
import pandas as pd
from datacheck.read_table import get_data_from_file
import numpy as np


def _get_filename(path):
    # 获取文件名
    if os.path.exists(path):
        return os.path.splitext(path)[0]
    return os.path.basename(path)


class CompareRowPlus:
    def __init__(self, file1, file2, file1_encoding='utf-8', file2_encoding='utf-8', save_path='result/result.xlsx',
                 directory='result', fields1=[], fields2=[], flag=1,
                 header=0):
        """
        :Description:  比较两个文件中指定字段是否相等，并输出对比结果到文件
        :flag=0,默认为0，比对输入相同结果，flag=1，比对输入不同结果，默认为1
        :param file1:  对比文件1
        :param file2:  对比文件2
        :param save_path: 对比后文件保存路径,默认result.xlsx，可自定义指定
        :param directory: 保存文件目录，默认为空，即当前目录
        :param encoding: 文件编码 默认utf-8
        :param header:   文件表头 默认0
        :param fields1:   指定字段，相当于唯一值，不能为空
        :param fields2:   指定字段，相当于唯一值，不能为空

        """
        self._flag = flag
        self._save_path = save_path
        self._directory = directory
        self._fields1 = fields1
        self._fields2 = fields2
        self._file1 = file1
        self._file2 = file2
        self._header = header
        # self._encoding = encoding  # 文件编码
        self._file1_encoding = file1_encoding  # 文件编码
        self._file2_encoding = file2_encoding  # 文件编码
        self._file_name1 = _get_filename(file1)  # 获取文件1名称
        self._file_name2 = _get_filename(file2)  # 获取文件2名称

    def _create_directory(self, directory):
        # 创建目录
        if self._directory:
            directory = self._directory
        if not os.path.exists(directory):
            os.makedirs(directory)
            return True
        return False

    def _file1_path(self):  # 获取文件1路径
        return self._file1

    def _file2_path(self):  # 获取文件2路径
        return self._file2

    def _read_file1(self, file_path, ):
        # 读取文件文件1
        if self._file1_encoding:
            return get_data_from_file(file_path, encoding=self._file1_encoding, header=self._header)
        else:
            return get_data_from_file(file_path)

    def _read_file2(self, file_path):
        # 读取文件文件2
        if self._file2_encoding:
            return get_data_from_file(file_path, encoding=self._file2_encoding, header=self._header)
        else:
            return get_data_from_file(file_path)

    def _save_file(self, common):
        # 另存为文件
        try:
            if self._save_path:
                self._save_path = self._save_path
            # self._set_background_color(common).to_excel(self._save_path, index=False, engine='openpyxl')
            common.to_excel(self._save_path)
            return True
        except PermissionError:
            print("file is open,please close it first")
            return False

    def _compare_data(self, df1, df2):
        """
        比较数据
        :param df1:
        :param df2:
        """
        def compare_and_style(fields1, fields2):
            # 确保数据框按指定字段排序
            df1_sorted = df1.sort_values(by=fields1).reset_index(drop=True)
            df2_sorted = df2.sort_values(by=fields2).reset_index(drop=True)

            # 创建布尔型DataFrame来标记不相同的数据
            diff_mask = pd.DataFrame(index=df1_sorted.index, columns=fields1)

            for i in range(len(fields1)):
                diff_mask[fields1[i]] = df1_sorted[fields1[i]] != df2_sorted[fields2[i]]

            # 将所有条件合并
            df1_sorted['is_equal'] = ~diff_mask.any(axis=1)
            common = df1_sorted[df1_sorted['is_equal']] if self._flag == 0 else df1_sorted[~df1_sorted['is_equal']]
            common = common.drop(columns=['is_equal'])

            # 定义样式函数
            def highlight_diff(data):
                attr = 'background-color: #d6a4a4'
                mask = diff_mask.reindex_like(data)
                return pd.DataFrame(np.where(mask, attr, ''), index=data.index, columns=data.columns)

            # 应用样式
            styled_common = common.style.apply(highlight_diff, axis=None)
            self._save_file(styled_common)
            print(
                f"{self._file_name1} 和 {self._file_name2} "
                f"compared done, {'same' if self._flag == 0 else 'different'} data save to {self._save_path}file")
            return styled_common

        if self._fields1 and self._fields2:
            return compare_and_style(self._fields1, self._fields2)
        else:
            raise ValueError("Fields must be specified for comparison")

    def run(self):
        df1 = self._read_file1(self._file1_path())  # 读取file1数据
        df2 = self._read_file2(self._file2_path())  # 读取file2数据

        if df2.equals(df1):
            print(f"{self._file_name1} 和 {self._file_name2} data is same about all fields")
            return ''
        if df1.shape[0] != df2.shape[0]:  # 数据条数不一致，不进行字段比较
            print(f"{self._file_name1} 和 {self._file_name2} data is not same about all fields")
            return ''
        self._create_directory(self._directory)  # 创建目录
        return self._compare_data(df1, df2)


if __name__ == '__main__':
    # 测试
    compare = CompareRowPlus('./files/fms_smt_bill_detail_alias_1729147283 - 副本.csv',
                             './files/fms_smt_bill_detail_alias_1729147283.csv',

                             flag=1,
                             file1_encoding='gbk',
                             file2_encoding='gbk',
                             fields1=["账号简称", "账号全称", "币种", "核算部门", "期初金额", "分账总额",
                                      "联盟佣金扣佣", "平台佣金扣佣", "分账退款总额", "联盟佣金退回", "佣金退回",
                                      "扣款总额"],
                             fields2=["账号简称", "账号全称", "币种", "核算部门", "期初金额", "分账总额",
                                      "联盟佣金扣佣", "平台佣金扣佣", "分账退款总额", "联盟佣金退回", "佣金退回",
                                      "扣款总额"],
                             directory="compare_result",
                             save_path='compare_result/result.xlsx',
                             header=0)

    compare.run()
